#include "Solver.hpp"

//smoothing function using Gauss-Seidel method
void Solver::smoothGS(Array &x, const Array &f, size_t iter)
{
    //assert(x.xsize() == f.xsize());
    
    size_t width = x.xsize();
    real h = x.hsize();
    //omega = 1.0 for Gauss-Seidel method
    real omg = 1.0;
    
    for ( size_t iter_ = 1 ; iter_<=iter ; ++iter_)
    {
        for (size_t i = 1; i < width-1; ++i)
        {
            for (size_t j = 1; j < width-1; ++j)
            {
                x(i,j) = (1-omg)*x(i,j) + omg*(4.0/h*h) * ( (x(i+1,j)+x(i-1,j+x(i,j+1)+x(i,j-1))/(h*h) -f(i,j) )
            }
        }
    }  
}

//smoothing function using red-black Gauss-Seidel
void Solver::smoothGSRB(Array &x, const Array &f, size_t iter)
{
    //assert(x.xsize() == f.xsize());
    
    size_t width = x.xsize();
    size_t j;
    real h = x.hSize();
    real omg = 1.0;
    
    //updating "red" cells
    for ( size_t iter_ = 1 ; iter_<=iter ; ++iter_)
    { 
        for (size_t i = 1; i < width-1; ++i)
        {
            if (i%2 == 1)
               j = 1;
            else
               j = 2;
            for (; j < width-1; j+=2)
            {
                x(i,j) = (1-omg)*x(i,j) + omg*(4.0/h*h) * ((x(i+1,j)+x(i-1,j+x(i,j+1)+x(i,j-1))/(h*h) -f(i,j))
            }
        }
    }
    
    //updating "black" cells
    for ( size_t iter_ = 1 ; iter_<=iter ; ++iter_)
    { 
        for (size_t i = 1; i < width-1; ++i)
        {
            if (i%2 == 1)
               j = 2;
            else
               j = 1;
            for (; j < width-1; j+=2)
            {
                x(i,j) = (1-omg)*x(i,j) + omg*(4.0/h*h) * ((x(i+1,j)+x(i-1,j+x(i,j+1)+x(i,j-1))/(h*h) -f(i,j))
            }
        }
    }
}

void Solver::restrict(const Array &r, Array &f)
{
    //          assert ( r. xsize () == r. ysize () );
    //        assert ( r. xsize () == 2*f. xsize () -1 );
    //      assert ( r. ysize () == 2*f. ysize () -1 );
    //    assert ( real (2)* r. hsize () == f. hsize () );

    const size_t size ( f. xsize () );
    for ( size_t i =1; i<size -1; ++i ) 
    {
        for ( size_t j =1; j<size -1; ++j )
        {
            f(i,j) = I_.C  * r(i*2   , j*2  ) +
                I_.N  * r(i *2+1 , j*2  ) +
                I_.S  * r(i*2 -1 , j*2  ) +
                I_.W  * r(i*2   , j *2 -1) +
                I_.E  * r(i*2   , j *2+1) +
                I_.NW * r(i *2+1 , j *2 -1) +
                I_.NE * r(i *2+1 , j *2+1) +
                I_.SW * r(i*2 -1 , j *2 -1) +
                I_.SE * r(i*2 -1 , j *2+1);
        }
    }
}

//Ax-f =r (the residual is not elevated to the second power)
void Solver::calcResidual(const Array &x, const Array &f, Array &r)
{
    const size_t size ( f. xsize () );
    for(size_t i=1;i<size-1;++i)
    {
        for(size_t j=1;j<size-1;++j)
        {
            r(i,j)=(x(i+1,j)-4.0*x(i,j)+x(i-1,j)+x(i,j+1)+x(i,j-1))/(x.hSize*x.hSize)-f(i,j);
        }
    }
}


void Solver::setStencil()
{
	
	A_.N = A_.S = A_.W = A_.E = 1;
	A_.NW =A_.NE=A_.SW = A_.SE =0;
	A_.C = -4;	

	I_.N = I_.S = I_.W = I_.E = 2/16;
	I_.NW =I_.NE=I_.SW = I_.SE =1/16;
	I_.C = 4/16;

	J_.N = J_.S = J_.W = J_.E = 2/4;
	J_.NW =J_.NE=J_.SW = J_.SE =1/4;
	J_.C = 4/4;

}
void Solver::prolongateAndCorrect(const Array &e, Array &x)
{
	//          assert ( r. xsize () == r. ysize () );
	//        assert ( r. xsize () == 2*f. xsize () -1 );
	//      assert ( r. ysize () == 2*f. ysize () -1 );
	//    assert ( real (2)* r. hsize () == f. hsize () );

	// Prolongate Equation: e[level-1] = J * e[level-2];
	// Correction Equation: x[level-1] = x[level-1] + e[level-1]			

	const size_t size ( e. xsize () );
	for ( size_t i =1; i<size -1; ++i ) 
	{
		for ( size_t j =1; j<size -1; ++j )
		{
			x(2*i,2*j) += e(i,j);
			x(2*i+1,2*j) += (e(i,j)+e(i+1,j))*J_.E;
			x(2*i-1,2*j) += (e(i,j)+e(i-1,j))*J_.W;
			x(2*i,2*j+1) += (e(i,j)+e(i,j+1))*J_.N;
			x(2*i,2*j-1) += (e(i,j)+e(i,j-1))*J_.S;
			x(2*i+1,2*j+1) += (e(i,j)+e(i,j+1)+e(i+1,j)+e(i+1,j+1))*J_.NE;
			x(2*i-1,2*j+1) += (e(i,j)+e(i,j+1)+e(i-1,j)+e(i-1,j+1))*J_.NW;
			x(2*i+1,2*j-1) += (e(i,j)+e(i,j-1)+e(i+1,j)+e(i+1,j-1))*J_.SE;
			x(2*i-1,2*j-1) += (e(i,j)+e(i,j-1)+e(i-1,j)+e(i-1,j-1))*J_.SW;
		}
	}
}
